// network/utils.go
package network

import (
	"gonum.org/v1/gonum/mat"
	"math"
)

// Sigmoid 激活函数
func Sigmoid(z *mat.VecDense) *mat.VecDense {
	out := mat.NewVecDense(z.Len(), nil)
	for i := 0; i < z.Len(); i++ {
		out.SetVec(i, 1/(1+math.Exp(-z.AtVec(i))))
	}
	return out
}

// SigmoidDerivative sigmoid的导数
func SigmoidDerivative(a *mat.VecDense) *mat.VecDense {
	out := mat.NewVecDense(a.Len(), nil)
	for i := 0; i < a.Len(); i++ {
		v := a.AtVec(i)
		out.SetVec(i, v*(1-v))
	}
	return out
}

// ReLU 激活函数
func ReLU(z *mat.VecDense) *mat.VecDense {
	out := mat.NewVecDense(z.Len(), nil)
	for i := 0; i < z.Len(); i++ {
		val := z.AtVec(i)
		if val > 0 {
			out.SetVec(i, val)
		} else {
			out.SetVec(i, 0)
		}
	}
	return out
}

// ReLUDerivative ReLU的导数
func ReLUDerivative(a *mat.VecDense) *mat.VecDense {
	out := mat.NewVecDense(a.Len(), nil)
	for i := 0; i < a.Len(); i++ {
		if a.AtVec(i) > 0 {
			out.SetVec(i, 1)
		} else {
			out.SetVec(i, 0)
		}
	}
	return out
}

// Softmax 函数
func Softmax(z *mat.VecDense) *mat.VecDense {
	sum := 0.0
	for i := 0; i < z.Len(); i++ {
		sum += math.Exp(z.AtVec(i))
	}
	out := mat.NewVecDense(z.Len(), nil)
	for i := 0; i < z.Len(); i++ {
		out.SetVec(i, math.Exp(z.AtVec(i))/sum)
	}
	return out
}

// 数据转换函数，用于将MNIST数据转换为CNN输入格式
func PrepareImageTensor(imageData []byte, height, width int) *Tensor3D {
	// 创建1通道的图像张量
	tensor := NewTensor3D(1, height, width)

	for h := 0; h < height; h++ {
		for w := 0; w < width; w++ {
			// 将像素值归一化到0-1之间
			pixel := float64(imageData[h*width+w]) / 255.0
			tensor.Data[0][h][w] = pixel
		}
	}

	return tensor
}

// 批量准备输入数据
func PrepareImageBatch(imageDataBatch [][]byte, height, width int) []interface{} {
	batchSize := len(imageDataBatch)
	batch := make([]interface{}, batchSize)

	for i := 0; i < batchSize; i++ {
		batch[i] = PrepareImageTensor(imageDataBatch[i], height, width)
	}

	return batch
}

// 准备目标数据（one-hot编码）
func PrepareTargets(labels []byte, numClasses int) []interface{} {
	numSamples := len(labels)
	targets := make([]interface{}, numSamples)

	for i := 0; i < numSamples; i++ {
		target := mat.NewVecDense(numClasses, nil)
		target.SetVec(int(labels[i]), 1.0)
		targets[i] = target
	}

	return targets
}
